
import os
import json
import time
import uuid
from datetime import datetime
from typing import List, Dict, Any
import asyncio
import logging

from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse, FileResponse
import pandas as pd
from PIL import Image
import pytesseract
from pdf2image import convert_from_path
import numpy as np

# 模拟DeepSeekOCR SDK调用
class DeepSeekOCRClient:
    """模拟DeepSeekOCR客户端"""
    
    @staticmethod
    async def recognize_text(image_path: str) -> Dict[str, Any]:
        """模拟OCR识别"""
        # 这里应该调用真实的DeepSeekOCR API
        # 为了演示，我们使用pytesseract作为模拟
        try:
            image = Image.open(image_path)
            text = pytesseract.image_to_string(image, lang='chi_sim+eng')
            
            return {
                "text": text,
                "confidence": 0.95,  # 模拟置信度
                "blocks": [
                    {
                        "text": line,
                        "confidence": 0.95,
                        "bbox": [0, 0, 100, 20]  # 模拟边界框
                    }
                    for line in text.split('\n') if line.strip()
                ]
            }
        except Exception as e:
            logging.error(f"OCR识别失败: {str(e)}")
            return {"text": "", "confidence": 0.0, "blocks": []}

app = FastAPI(title="DeepSeekOCR-ESG测试平台", version="1.0.0")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 挂载静态文件
app.mount("/static", StaticFiles(directory="static"), name="static")

# 全局变量存储测试数据
test_results = {}
ground_truth_data = {}

@app.get("/")
async def root():
    return {"message": "DeepSeekOCR-ESG测试平台API服务"}

@app.post("/api/upload-report")
async def upload_report(file: UploadFile = File(...)):
    """上传ESG报告"""
    try:
        # 生成唯一ID
        report_id = str(uuid.uuid4())
        
        # 保存文件
        file_extension = file.filename.split('.')[-1].lower()
        file_path = f"data/reports/{report_id}.{file_extension}"
        
        os.makedirs("data/reports", exist_ok=True)
        with open(file_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)
        
        # 处理PDF或图片
        images = []
        if file_extension == 'pdf':
            # 转换PDF为图片
            pages = convert_from_path(file_path, dpi=300)
            for i, page in enumerate(pages):
                img_path = f"data/reports/{report_id}_page_{i+1}.png"
                page.save(img_path, 'PNG')
                images.append(img_path)
        else:
            # 直接保存图片
            images = [file_path]
        
        return {
            "report_id": report_id,
            "filename": file.filename,
            "pages": len(images),
            "images": images
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/perform-ocr")
async def perform_ocr(report_id: str = Form(...)):
    """对报告执行OCR识别"""
    try:
        # 获取报告图片
        report_files = [f for f in os.listdir("data/reports") if f.startswith(report_id)]
        image_files = [f for f in report_files if f.endswith(('.png', '.jpg', '.jpeg'))]
        
        ocr_results = {}
        total_time = 0
        
        for img_file in sorted(image_files):
            img_path = f"data/reports/{img_file}"
            start_time = time.time()
            
            # 调用OCR
            result = await DeepSeekOCRClient.recognize_text(img_path)
            
            end_time = time.time()
            processing_time = end_time - start_time
            total_time += processing_time
            
            ocr_results[img_file] = {
                "text": result["text"],
                "confidence": result["confidence"],
                "processing_time": processing_time,
                "blocks": result["blocks"]
            }
        
        # 保存结果
        test_results[report_id] = {
            "ocr_results": ocr_results,
            "total_time": total_time,
            "timestamp": datetime.now().isoformat()
        }
        
        return {
            "report_id": report_id,
            "results": ocr_results,
            "total_processing_time": total_time
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/upload-ground-truth")
async def upload_ground_truth(
    report_id: str = Form(...),
    ground_truth: UploadFile = File(...)
):
    """上传人工标注的标准答案"""
    try:
        content = await ground_truth.read()
        text = content.decode('utf-8')
        
        # 按页分割标注数据
        pages = text.split('---PAGE_BREAK---')
        ground_truth_data[report_id] = {
            "pages": pages,
            "timestamp": datetime.now().isoformat()
        }
        
        return {"message": "标准答案上传成功", "pages": len(pages)}
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/calculate-accuracy")
async def calculate_accuracy(report_id: str = Form(...)):
    """计算准确率"""
    try:
        if report_id not in test_results:
            raise HTTPException(status_code=404, detail="未找到OCR结果")
        
        if report_id not in ground_truth_data:
            raise HTTPException(status_code=404, detail="未找到标准答案")
        
        ocr_results = test_results[report_id]["ocr_results"]
        ground_truth_pages = ground_truth_data[report_id]["pages"]
        
        accuracy_results = {}
        
        for img_file, ocr_data in ocr_results.items():
            ocr_text = ocr_data["text"]
            
            # 找到对应的标注页
            page_num = int(img_file.split('_')[-1].split('.')[0]) - 1
            if page_num < len(ground_truth_pages):
                ground_truth_text = ground_truth_pages[page_num]
                
                # 计算字符级准确率
                char_accuracy = calculate_character_accuracy(ocr_text, ground_truth_text)
                
                # 计算段落级准确率
                para_accuracy = calculate_paragraph_accuracy(ocr_text, ground_truth_text)
                
                # 计算关键字段召回率
                field_recall = calculate_field_recall(ocr_text, ground_truth_text)
                
                accuracy_results[img_file] = {
                    "character_accuracy": char_accuracy,
                    "paragraph_accuracy": para_accuracy,
                    "field_recall": field_recall,
                    "ocr_text_length": len(ocr_text),
                    "ground_truth_length": len(ground_truth_text)
                }
        
        # 计算总体准确率
        overall_accuracy = {
            "character_accuracy": np.mean([r["character_accuracy"] for r in accuracy_results.values()]),
            "paragraph_accuracy": np.mean([r["paragraph_accuracy"] for r in accuracy_results.values()]),
            "field_recall": np.mean([r["field_recall"] for r in accuracy_results.values()]),
            "total_pages": len(accuracy_results)
        }
        
        # 保存结果
        result_file = f"data/results/{report_id}_accuracy.json"
        os.makedirs("data/results", exist_ok=True)
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump({
                "report_id": report_id,
                "overall_accuracy": overall_accuracy,
                "page_results": accuracy_results,
                "timestamp": datetime.now().isoformat()
            }, f, ensure_ascii=False, indent=2)
        
        return {
            "report_id": report_id,
            "overall_accuracy": overall_accuracy,
            "page_results": accuracy_results
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/download-report/{report_id}")
async def download_report(report_id: str):
    """下载测试报告"""
    try:
        result_file = f"data/results/{report_id}_accuracy.json"
        if not os.path.exists(result_file):
            raise HTTPException(status_code=404, detail="测试报告不存在")
        
        return FileResponse(
            result_file,
            media_type='application/json',
            filename=f"esg_ocr_test_report_{report_id}.json"
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def calculate_character_accuracy(ocr_text: str, ground_truth: str) -> float:
    """计算字符级准确率"""
    ocr_chars = list(ocr_text.replace(' ', '').replace('\n', ''))
    truth_chars = list(ground_truth.replace(' ', '').replace('\n', ''))
    
    # 使用编辑距离计算准确率
    from difflib import SequenceMatcher
    return SequenceMatcher(None, ''.join(ocr_chars), ''.join(truth_chars)).ratio()

def calculate_paragraph_accuracy(ocr_text: str, ground_truth: str) -> float:
    """计算段落级准确率"""
    ocr_paragraphs = [p.strip() for p in ocr_text.split('\n\n') if p.strip()]
    truth_paragraphs = [p.strip() for p in ground_truth.split('\n\n') if p.strip()]
    
    if not truth_paragraphs:
        return 1.0
    
    matched = 0
    for truth_para in truth_paragraphs:
        for ocr_para in ocr_paragraphs:
            if truth_para in ocr_para or ocr_para in truth_para:
                matched += 1
                break
    
    return matched / len(truth_paragraphs)

def calculate_field_recall(ocr_text: str, ground_truth: str) -> float:
    """计算关键字段召回率"""
    # 定义ESG关键字段
    esg_fields = [
        "碳排放", "温室气体", "能源消耗", "水资源", "废物处理",
        "员工人数", "安全事故", "培训时长", "董事会", "治理结构"
    ]
    
    found_fields = 0
    total_fields = 0
    
    for field in esg_fields:
        if field in ground_truth:
            total_fields += 1
            if field in ocr_text:
                found_fields += 1
    
    return found_fields / total_fields if total_fields > 0 else 0.0

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
